#include "kernel_operator.h"

using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2;

class KernelAsinhGradFloat
{
    using T = float;

public:
    __aicore__ inline KernelAsinhGradFloat() {}
    __aicore__ inline void Init(GM_ADDR y, GM_ADDR dy, GM_ADDR z,
                                uint32_t coreDataNum, uint32_t finalTileNum, uint32_t tileDataNum, uint32_t tailDataNum, TPipe *pipeIn)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->coreDataNum = coreDataNum;
        this->tileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->tailDataNum = tailDataNum;

        yGm.SetGlobalBuffer((__gm__ T *)y, this->coreDataNum);
        dyGm.SetGlobalBuffer((__gm__ T *)dy, this->coreDataNum);
        zGm.SetGlobalBuffer((__gm__ T *)z, this->coreDataNum);

        pipe = pipeIn;
        pipe->InitBuffer(inQueueY, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe->InitBuffer(inQueueDy, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe->InitBuffer(outQueueZ, BUFFER_NUM, this->tileDataNum * sizeof(T));
    }
    __aicore__ inline void Process()
    {
        uint32_t loopCount = this->tileNum;
        // 每次处理的数据量，最后一次是 使用尾块
        this->processDataNum = this->tileDataNum;
        for (uint32_t i = 0; i < loopCount; i++)
        {
            if (i == this->tileNum - 1)
            {
                this->processDataNum = this->tailDataNum;
            }
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(uint32_t progress)
    {
        LocalTensor<T> yLocal = inQueueY.AllocTensor<T>();
        LocalTensor<T> dyLocal = inQueueDy.AllocTensor<T>();
        DataCopy(yLocal, yGm[progress * this->tileDataNum], this->processDataNum);
        DataCopy(dyLocal, dyGm[progress * this->tileDataNum], this->processDataNum);
        inQueueY.EnQue(yLocal);
        inQueueDy.EnQue(dyLocal);
    }
    __aicore__ inline void Compute(uint32_t progress)
    {
        LocalTensor<T> yLocal = inQueueY.DeQue<T>();
        LocalTensor<T> dyLocal = inQueueDy.DeQue<T>();
        LocalTensor<T> zLocal = outQueueZ.AllocTensor<T>();

        // 2dy
        Add(dyLocal, dyLocal, dyLocal, this->processDataNum);
        // Muls(dyLocal, dyLocal, static_cast<T>(2), this->processDataNum);

        // exp(y)
        Exp(zLocal, yLocal, this->processDataNum);
        Muls(yLocal, yLocal, static_cast<T>(-1), this->processDataNum);
        Exp(yLocal, yLocal, this->processDataNum);
        // Reciprocal(zLocal, yLocal, this->processDataNum);
        Add(zLocal, zLocal, yLocal, this->processDataNum);

        Div(zLocal, dyLocal, zLocal, this->processDataNum);

        outQueueZ.EnQue(zLocal);
        inQueueY.FreeTensor(yLocal);
        inQueueDy.FreeTensor(dyLocal);
    }
    __aicore__ inline void CopyOut(uint32_t progress)
    {
        LocalTensor<T> zLocal = outQueueZ.DeQue<T>();
        DataCopy(zGm[progress * this->tileDataNum], zLocal, this->processDataNum);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    TPipe *pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueY, inQueueDy;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueZ;

    GlobalTensor<T> yGm;
    GlobalTensor<T> dyGm;
    GlobalTensor<T> zGm;

    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tileDataNum;
    uint32_t tailDataNum;
    uint32_t processDataNum;
};

class KernelAsinhGradHalf
{
    using T = half;
    using Tfp32 = float;

public:
    __aicore__ inline KernelAsinhGradHalf() {}
    __aicore__ inline void Init(GM_ADDR y, GM_ADDR dy, GM_ADDR z,
                                uint32_t coreDataNum, uint32_t finalTileNum, uint32_t tileDataNum, uint32_t tailDataNum, TPipe *pipeIn)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->coreDataNum = coreDataNum;
        this->tileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->tailDataNum = tailDataNum;

        yGm.SetGlobalBuffer((__gm__ T *)y, this->coreDataNum);
        dyGm.SetGlobalBuffer((__gm__ T *)dy, this->coreDataNum);
        zGm.SetGlobalBuffer((__gm__ T *)z, this->coreDataNum);
        pipe = pipeIn;

        pipe->InitBuffer(inQueueY, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe->InitBuffer(inQueueDy, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe->InitBuffer(outQueueZ, BUFFER_NUM, this->tileDataNum * sizeof(T));

        pipe->InitBuffer(tmpBuffer, 3 * this->tileDataNum * sizeof(Tfp32));
    }
    __aicore__ inline void Process()
    {
        uint32_t loopCount = this->tileNum;
        // 每次处理的数据量，最后一次是 使用尾块
        this->processDataNum = this->tileDataNum;
        for (uint32_t i = 0; i < loopCount; i++)
        {
            if (i == this->tileNum - 1)
            {
                this->processDataNum = this->tailDataNum;
            }
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }
    __aicore__ inline void Process2()
    {
        uint32_t loopCount = (this->tileNum - 1) * this->tileDataNum + this->tailDataNum;
        LocalTensor<float> yLocal = inQueueY.AllocTensor<float>();
        float tmp = 3;
        for (uint32_t i = 0; i < loopCount; i++)
        {
            float dy = static_cast<float>(dyGm.GetValue(i));
            float y = static_cast<float>(yGm.GetValue(i));
            y = (y < 0) ? -y : y;
            dy += dy;

            if (y < tmp)
            {
                // 两次
                yLocal.SetValue(0, y);
                Exp(yLocal, yLocal, 1);
                y = yLocal.GetValue(0) + 1 / yLocal.GetValue(0);
                zGm.SetValue(i, static_cast<T>(dy / y));
            }
            else
            {
                // 一次
                yLocal.SetValue(0, y);
                Exp(yLocal, yLocal, 1);
                zGm.SetValue(i, static_cast<T>(dy / yLocal.GetValue(0)));
            }
        }
        inQueueY.FreeTensor(yLocal);
    }

private:
    __aicore__ inline void CopyIn(uint32_t progress)
    {
        LocalTensor<T> yLocal = inQueueY.AllocTensor<T>();
        LocalTensor<T> dyLocal = inQueueDy.AllocTensor<T>();
        DataCopy(yLocal, yGm[progress * this->tileDataNum], this->processDataNum);
        DataCopy(dyLocal, dyGm[progress * this->tileDataNum], this->processDataNum);
        inQueueY.EnQue(yLocal);
        inQueueDy.EnQue(dyLocal);
    }
    __aicore__ inline void Compute(uint32_t progress)
    {
        LocalTensor<T> yLocal = inQueueY.DeQue<T>();
        LocalTensor<T> dyLocal = inQueueDy.DeQue<T>();
        LocalTensor<T> zLocal = outQueueZ.AllocTensor<T>();
        auto yBuf = tmpBuffer.Get<Tfp32>(this->processDataNum);
        auto zBuf = tmpBuffer.GetWithOffset<Tfp32>(this->processDataNum, this->processDataNum * sizeof(Tfp32));

        Add(dyLocal, dyLocal, dyLocal, this->processDataNum);
        // Muls(dyLocal, dyLocal, static_cast<T>(2), this->processDataNum);

        Cast(yBuf, yLocal, RoundMode::CAST_NONE, this->processDataNum);
        // Exp(zBuf, yBuf, this->processDataNum);
        // Muls(yBuf, yBuf, static_cast<Tfp32>(-1), this->processDataNum);
        // Exp(yBuf, yBuf, this->processDataNum);
        Exp(yBuf, yBuf, this->processDataNum);
        Duplicate(zBuf, static_cast<Tfp32>(1), this->processDataNum);
        Div(zBuf, zBuf, yBuf, this->processDataNum);

        Add(zBuf, zBuf, yBuf, this->processDataNum);

        Cast(zLocal, zBuf, RoundMode::CAST_NONE, this->processDataNum);
        Div(zLocal, dyLocal, zLocal, this->processDataNum);

        outQueueZ.EnQue(zLocal);
        inQueueY.FreeTensor(yLocal);
        inQueueDy.FreeTensor(dyLocal);
    }
    __aicore__ inline void Compute2(uint32_t progress)
    {
        LocalTensor<T> yLocal = inQueueY.DeQue<T>();
        LocalTensor<T> dyLocal = inQueueDy.DeQue<T>();
        LocalTensor<T> zLocal = outQueueZ.AllocTensor<T>();
        auto yBuf = tmpBuffer.Get<Tfp32>(this->processDataNum);
        auto zBuf = tmpBuffer.GetWithOffset<Tfp32>(this->processDataNum, this->processDataNum * sizeof(Tfp32));
        auto tmp = tmpBuffer.GetWithOffset<Tfp32>(this->processDataNum, 2 * this->processDataNum * sizeof(Tfp32));

        Cast(yBuf, yLocal, RoundMode::CAST_NONE, this->processDataNum);
        Mul(tmp, yBuf, yBuf, this->processDataNum);                      // x^2
        Muls(tmp, tmp, static_cast<Tfp32>(1 / 2), this->processDataNum); // x^2 / 2
        Adds(zBuf, tmp, static_cast<Tfp32>(1), this->processDataNum);    // 1 + x^2 / 2!

        Mul(yBuf, tmp, tmp, this->processDataNum);                         // x^4 / (2 * 2)
        Muls(yBuf, yBuf, static_cast<Tfp32>(1 / 6), this->processDataNum); // x^4 / 4!
        Add(zBuf, yBuf, zBuf, this->processDataNum);                       // 1 + x^2 / 2! + x^4 / 4!

        Muls(yBuf, yBuf, static_cast<Tfp32>(1 / 15), this->processDataNum); // x^6 / 6!
        Mul(yBuf, tmp, yBuf, this->processDataNum);                         // x^6 / (2 * 4!)
        Add(zBuf, yBuf, zBuf, this->processDataNum);                        // 1 + x^2 / 2! + x^4 / 4! + x^6 / 6!

        Muls(yBuf, yBuf, static_cast<Tfp32>(1 / 28), this->processDataNum); // x^8 / 8!
        Mul(yBuf, tmp, yBuf, this->processDataNum);                         // x^8 / (2 * 6!)
        Add(zBuf, yBuf, zBuf, this->processDataNum);                        // 1 + x^2 / 2! + x^4 / 4! + x^6 / 6! + x^8 / 8!
        Cast(tmp, dyLocal, RoundMode::CAST_NONE, this->processDataNum);
        Div(zBuf, tmp, zBuf, this->processDataNum);

        Cast(zLocal, zBuf, RoundMode::CAST_NONE, this->processDataNum);
        // Div(zLocal, dyLocal, zLocal, this->processDataNum);

        outQueueZ.EnQue(zLocal);
        inQueueY.FreeTensor(yLocal);
        inQueueDy.FreeTensor(dyLocal);
    }

    __aicore__ inline void CopyOut(uint32_t progress)
    {
        LocalTensor<T> zLocal = outQueueZ.DeQue<T>();
        DataCopy(zGm[progress * this->tileDataNum], zLocal, this->processDataNum);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    TPipe *pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueY, inQueueDy;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
    TBuf<QuePosition::VECCALC> tmpBuffer;

    GlobalTensor<T> yGm;
    GlobalTensor<T> dyGm;
    GlobalTensor<T> zGm;

    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tileDataNum;
    uint32_t tailDataNum;
    uint32_t processDataNum;
};
extern "C" __global__ __aicore__ void asinh_grad(GM_ADDR y, GM_ADDR dy, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);

    TPipe pipe;
    if (TILING_KEY_IS(2))
    {
        KernelAsinhGradHalf op;
        op.Init(y, dy, z,
                tiling_data.coreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.tailDataNum, &pipe);
        op.Process();
    }
    else if (TILING_KEY_IS(1))
    {
        KernelAsinhGradFloat op;
        op.Init(y, dy, z,
                tiling_data.coreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.tailDataNum, &pipe);
        op.Process();
    }
}